package org.example.myleet.p437;

import org.example.myleet.Utils.TreeNode;

import java.util.HashMap;

/**
 * 更优的方法，在树状结构中实现前缀和
 * 每遍历到下一个节点的时候计算到达当前节点cur的前缀和curPreSum，并记忆化出现这种前缀和的路径数（例如HashMap），方便查找
 * 然后根据curPreSum - targetSum的结果（相当于从根节点到某个A节点的前缀和，然后A节点到cur节点的前缀和等于targetSum）
 * 然后在记忆化中搜寻是否出现过这样的前缀，如果出现次数大于0，路径计数加上出现次数
 * 注意⚠️
 * 1、初始化时需要加入preSum为0，路径数为1的值，因为什么都不选的时候也算一种可能情况（不是结果），在后面查找前缀和的时候有用
 * 2、需要在当前节点搜索结束时进行回溯，在记忆化前缀和中去当前节点的前缀和，因为父节点换右边子节点搜索的时候，左子节点的前缀和失效了
 */
public class Solution {

    private int count;

    private int targetSum;

    private HashMap<Integer, Integer> preSumCountMap;

    public int pathSum(TreeNode root, int targetSum) {
        if (root == null) {
            return 0;
        }
        preSumCountMap = new HashMap<>();
        preSumCountMap.put(0, 1);
        count = 0;
        this.targetSum = targetSum;
        dfs(root, 0);
        return count;
    }

    /**
     * @param node 当前节点
     * @param preSum 父节点的前缀和
     */
    private void dfs(TreeNode node, int preSum) {
        //求出当前节点的前缀和
        int curPreSum = node.val + preSum;
        //计算curPreSum - targetSum，即根节点到某个cur的祖先节点的前缀和，并从记忆化的前缀和中捞取
        int pathCount = preSumCountMap.getOrDefault(curPreSum - targetSum, 0);
        //路径计数累计
        count += pathCount;
        //当前前缀和的路径数+1，并记忆化
        int newSumCount = preSumCountMap.getOrDefault(curPreSum, 0);
        ++newSumCount;
        preSumCountMap.put(curPreSum, newSumCount);
        //搜索左右分支
        if (node.left != null) {
            dfs(node.left, curPreSum);
        }
        if (node.right != null) {
            dfs(node.right, curPreSum);
        }
        //回溯，如果前缀和路径数为0，则删除对应记忆
        if (newSumCount < 2) {
            preSumCountMap.remove(curPreSum);
        } else {
            preSumCountMap.put(curPreSum, newSumCount - 1);
        }
    }
}
